import argparse
import os
import numpy as np
import pandas as pd
from tqdm import tqdm
from sklearn.metrics import ndcg_score

def z_score_norm(x):
    return (x - x.mean()) / (x.std())

def ndcg_at_k(y_pred, y_true, k=None):
    y_pred, y_true = np.array(y_pred), np.array(y_true)
    return ndcg_score(z_score_norm(y_true).reshape(1, -1), z_score_norm(y_pred).reshape(1, -1), k=k)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--result_path', type=str, default='result/tm-benchmark')
    parser.add_argument('--pred_score', type=str, default='k20_h1280')
    parser.add_argument('--true_score', type=str, default='score')
    parser.add_argument('--benchmark_path', type=str, default='data/tm-benchmark/DATASET')
    parser.add_argument('--k', type=int, default=5)
    args = parser.parse_args()
    
    ndcg_res = {args.pred_score: []}
    proteins = os.listdir(args.result_path)
    proteins = [p for p in proteins if os.path.isdir(os.path.join(args.result_path, p))]
    for p in proteins:
        df_score = pd.read_table(os.path.join(args.result_path, p, p + "_labeled.tsv"))
        ndcg = ndcg(df_score['pred_score'], df_score['true_score'], args.k)
        ndcg_res[args.pred_score].append(ndcg)
    
    proteins = os.listdir(args.benchmark_path)
    proteins = [p for p in proteins if os.path.isdir(os.path.join(args.benchmark_path, p))]
    for p in tqdm(proteins):
        benchmark_models = os.listdir(os.path.join(args.benchmark_path, p, "predictions"))
        benchmark_models = [m.split(".")[-2] for m in benchmark_models if m.endswith(".tsv")]
        for m in benchmark_models:
            if ndcg_res.get(m) is None:
                ndcg_res[m] = []
            df_score = pd.read_table(os.path.join(args.benchmark_path, p, "predictions", f"{p}.{m}.tsv"))
            pred_score = df_score['score']
            true_score = pd.read_table(os.path.join(args.benchmark_path, p, "experiments", f"{p}.tsv"))['score']
            ndcg = ndcg_at_k(pred_score, true_score, args.k)
            ndcg_res[m].append(ndcg)
            
    
    res_df = pd.DataFrame(ndcg_res)
    res_df.to_csv(os.path.join(args.result_path, f"ndcg_k{args.k}.csv"), index=False)
